#from geesibling.core.lib._graph import Graph as RawGraph

import traceback
import torch
#from torch import nn
#from torch import fx
from torch.fx import symbolic_trace
from torchvision import models


class Get_IR(torch.fx.Interpreter):
    def run_node(self, n):
        try:
            result = super().run_node(n)
        except Exception:
            traceback.print_exc()
            raise RuntimeError(f"Error while run node:{n.format_node()}")

        is_find = False

        def extract_meta(t):
            if isinstance(t, torch.Tensor):
                nonlocal is_find
                is_find = True
                return _extra_meta(t)
            else:
                return t

        def _extra_meta(t):
            if n.op == "call_module":
                submod = self.fetch_attr(n.target)
                return {'name': n.name, 'op': n.op, 'args': n.args, 'outputShape': t.shape, 'kw': n.kwargs, 'mod': submod}
            elif n.op == "call_method":
                return {'name': n.name, 'op': n.op, 'args': n.args, 'outputShape': t.shape, 'kw': n.kwargs}
            elif n.op == "call_function":
                return {'name': n.name, 'op': n.op, 'args': n.args, 'outputShape': t.shape, 'kw': n.kwargs}
            else:
                return {'name': n.name, 'op': n.op, 'args': n.args, 'outputShape': t.shape,}

        n.meta["result"] = torch.fx.node.map_aggregate(result, extract_meta)
        n.meta["find"] = is_find
        return result


# class MyModule(torch.nn.Module):
#     def __init__(self):
#         super().__init__()
#         self.param = torch.nn.Parameter(torch.rand(3, 4))
#         self.linear = torch.nn.Linear(4, 5)

#     def forward(self, x):
#         return torch.topk(torch.sum(
#             self.linear(x + self.linear.weight).relu(), dim=-1), 3)

# model = MyModule()
model = models.resnet18(pretrained = False)

traced = symbolic_trace(model)
input_example = torch.randn((1, 64, 3, 224, 224))
# print(model(input_example))

args = (input_example)
kwargs = {}
_ = Get_IR(traced).run(*args, **kwargs)
print(traced.graph.print_tabular())
for node in traced.graph.nodes:
    input_shape = []
    for i in node._input_nodes:
        if isinstance(i.meta['result'],tuple):
            input_shape.append(i.meta['result'][0]['outputShape'])
        else:
            input_shape.append(i.meta['result']['outputShape'])
    if isinstance(node.meta['result'],dict):
        node.meta['result']['inputShape'] = input_shape
    else:
        node.meta['result'] = node.meta['result'][0]
        node.meta['result']['intputShape']  = input_shape
    # node._pretty_print_target(node.target)
    print(node.meta)
# print(traced(input_example))


# traced.to_folder(r'G:\idea program\pytorch\pytorch\torch\fx\examples\resnet18')

#graph = RawGraph()
#print(graph)
